package main.java.lpr;

import java.awt.image.BufferedImage;
import java.nio.file.Paths;
import java.util.ArrayList;
import java.util.Arrays;
import java.util.List;

import org.bytedeco.javacv.Frame;
import org.bytedeco.javacv.Java2DFrameConverter;
import org.bytedeco.javacv.OpenCVFrameConverter;
import org.bytedeco.opencv.opencv_core.Mat;
import org.bytedeco.opencv.opencv_core.Point2f;
import org.bytedeco.opencv.opencv_core.Size;

import ai.djl.inference.Predictor;
import ai.djl.modality.cv.Image;
import ai.djl.modality.cv.ImageFactory;
import ai.djl.modality.cv.output.BoundingBox;
import ai.djl.modality.cv.output.DetectedObjects;
import ai.djl.modality.cv.output.Rectangle;
import ai.djl.modality.cv.output.DetectedObjects.DetectedObject;
import ai.djl.modality.cv.translator.YoloV5Translator;
import ai.djl.modality.cv.util.NDImageUtils;
import ai.djl.repository.zoo.Criteria;
import ai.djl.repository.zoo.ZooModel;
import ai.djl.translate.Translator;
import ai.djl.training.util.ProgressBar;

import static org.bytedeco.opencv.global.opencv_imgproc.resize;
import static org.bytedeco.opencv.global.opencv_imgproc.getRectSubPix;

/**
 * LPR
 * 
 * @date: 2022-02-28 13:44
 * @author: alvin
 * @copyright: Copyright © 2022 YZZH. All rights reserved.
 * @description: main.java.lpr/LPR.java
 */
public class LPR {

    private String areaModelPath = "models/lpi.torchscript.pt";

    private String charModelPath = "models/lpr_gru";

    private ZooModel<Image, DetectedObjects> areaModel = null;

    private ZooModel<Image, PlateChar> charModel = null;

    private static final List<String> CLASSES = Arrays.asList(
            "京", "沪", "津", "渝", "冀", "晋", "蒙", "辽", "吉", "黑", "苏", "浙", "皖", "闽", "赣", "鲁", "豫", "鄂", "湘", "粤", "桂",
            "琼", "川", "贵", "云", "藏", "陕", "甘", "青", "宁", "新",
            "0", "1", "2", "3", "4", "5", "6", "7", "8", "9",
            "A", "B", "C", "D", "E", "F", "G", "H", "J", "K", "L", "M", "N", "P", "Q", "R", "S", "T", "U", "V", "W",
            "X", "Y", "Z",
            "港", "学", "使", "警", "澳", "挂", "军", "北", "南", "广", "沈", "兰", "成", "济", "海", "民", "航", "空");

    public LPR() {
        loadAreaModel();
        loadCharModel();
    }

    private void loadAreaModel() {
        Translator<Image, DetectedObjects> translator = YoloV5Translator.builder()
                .optSynsetArtifactName("lpc.names")
                .build();
        Criteria<Image, DetectedObjects> criteria = Criteria.builder()
                .setTypes(Image.class, DetectedObjects.class)
                .optModelPath(Paths.get(areaModelPath))
                .optModelName("lpi.torchscript.pt")
                .optTranslator(translator)
                .optEngine("PyTorch")
                .optProgress(new ProgressBar())
                .build();
        try {
            areaModel = criteria.loadModel();
        } catch (Exception e) {
            e.printStackTrace();
        }
    }

    private void loadCharModel() {
        LprE2eTranslator translator = LprE2eTranslator.builder()
                .addTransform(a -> NDImageUtils.resize(a, 164, 48).transpose(1, 0, 2))
                .optClasses(CLASSES)
                .build();
        Criteria<Image, PlateChar> criteria = Criteria.builder()
                .setTypes(Image.class, PlateChar.class)
                .optModelPath(Paths.get(charModelPath))
                .optEngine("TensorFlow")
                .optModelName("saved_model")
                .optTranslator(translator)
                .optProgress(new ProgressBar())
                .build();
        try {
            charModel = criteria.loadModel();
        } catch (Exception e) {
            e.printStackTrace();
        }
    }

    public List<Plate> detection(Mat img, double areaThresh, double charThresh) {
        List<Plate> plates = detectionPlateAreas(img, areaThresh);
        for (Plate plate : plates) {
            Mat plateImg = plate.getMat();
            List<PlateChar> plateChars = detectionPlateChars(plateImg, charThresh);
            if (!plateChars.isEmpty()) {
                plate.setPlateResult(plateChars.get(0));
            }
        }
        return plates;
    }

    public List<Plate> detectionPlateAreas(Mat img, double areaThresh) {
        List<Plate> plates = new ArrayList<Plate>();
        if (img == null || img.empty()) {
            return plates;
        }

        Mat src = new Mat();
        Size size = new Size(480, 480);
        resize(img, src, size);
        size.close();

        // 转换成BufferedImage
        OpenCVFrameConverter.ToMat matConv = new OpenCVFrameConverter.ToMat();
        Java2DFrameConverter biConv = new Java2DFrameConverter();
        Frame frame = matConv.convert(src);
        BufferedImage bimg = biConv.getBufferedImage(frame);

        // 释放
        frame.close();
        biConv.close();
        matConv.close();

        // 转换Image
        Image djl_img = ImageFactory.getInstance().fromImage(bimg);

        // 构建预测者
        Predictor<Image, DetectedObjects> predictor = areaModel.newPredictor();

        try {
            // 检测出全部对象
            DetectedObjects objects = predictor.predict(djl_img);

            for (DetectedObject obj : objects.<DetectedObject>items()) {
                if (areaThresh < obj.getProbability()) {
                    BoundingBox bbox = obj.getBoundingBox();
                    Rectangle rectangle = bbox.getBounds();

                    int h = (int) (rectangle.getHeight() * img.arrayHeight() / 480);
                    int w = (int) (rectangle.getWidth() * img.arrayWidth() / 480);
                    int x = (int) (rectangle.getX() * img.arrayWidth() / 480);
                    int y = (int) (rectangle.getY() * img.arrayHeight() / 480);

                    // coco对象
                    String type = "";
                    try {
                        type = obj.getClassName().trim();
                    } catch (Exception e) {

                    }
                    Plate o = new Plate();
                    o.setType(type);
                    o.setConfidence((float) obj.getProbability());
                    o.setH(h + 10);
                    o.setW(w + 10);
                    o.setX(x - 5);// 左上角的点
                    o.setY(y - 5);// 左上角的点

                    // 截取对象
                    Mat mat = new Mat();
                    Size si = new Size(o.getW(), o.getH());
                    Point2f p = new Point2f(o.getX() + o.getW() / 2, o.getY() + o.getH() / 2);
                    getRectSubPix(img, si, p, mat);
                    o.setMat(mat);

                    p.close();
                    si.close();

                    // 保存对象
                    plates.add(o);
                }
            }
        } catch (Exception e) {
            e.printStackTrace();
        }
        src.close();
        predictor.close();
        return plates;
    }

    public List<PlateChar> detectionPlateChars(Mat img, double thresh) {
        List<PlateChar> plates = new ArrayList<PlateChar>();
        if (img == null || img.empty()) {
            return plates;
        }

        // 转换成BufferedImage
        OpenCVFrameConverter.ToMat matConv = new OpenCVFrameConverter.ToMat();
        Java2DFrameConverter biConv = new Java2DFrameConverter();
        Frame frame = matConv.convert(img);
        BufferedImage bimg = biConv.getBufferedImage(frame);

        // 释放
        frame.close();
        biConv.close();
        matConv.close();

        // 转换Image
        Image djl_img = ImageFactory.getInstance().fromImage(bimg);

        // 构建预测者
        Predictor<Image, PlateChar> predictor = charModel.newPredictor();

        try {
            // 检测出全部对象
            PlateChar plateChar = predictor.predict(djl_img);
            if (plateChar.getConfidence() > thresh) {
                plates.add(plateChar);
            }
        } catch (Exception e) {
            e.printStackTrace();
        }
        predictor.close();
        return plates;
    }
}
